import torch
import torch_geometric.graphgym.register as register
from torch_geometric.graphgym.config import cfg
from torch_geometric.graphgym.models.gnn import GNNPreMP
from torch_geometric.graphgym.models.layer import (new_layer_config,
                                                   BatchNorm1dNode)
from torch_geometric.graphgym.register import register_network
from graphgps.encoder.ER_edge_encoder import EREdgeEncoder

from graphgps.layer.gps_layer import GPSLayer


class FeatureEncoder(torch.nn.Module):
    """
    Encoding node and edge features

    特征维度： dim_in → cfg.gnn.dim_inner
    Args:
        dim_in (int): Input feature dimension
    """
    def __init__(self, dim_in):
        super(FeatureEncoder, self).__init__()
        self.dim_in = dim_in
        # 是否使用节点编码器
        if cfg.dataset.node_encoder:
            # Encode integer node features via nn.Embeddings
            # 动态选择边编码器类型
            NodeEncoder = register.node_encoder_dict[
                cfg.dataset.node_encoder_name]
            # 创建节点编码器实例，输入特征维度为 cfg.gnn.dim_inner。
            self.node_encoder = NodeEncoder(cfg.gnn.dim_inner)
            # 根据配置是否使用批量归一化 (BatchNorm1dNode)，可以为边或节点的编码器添加批量归一化层。
            if cfg.dataset.node_encoder_bn:
                self.node_encoder_bn = BatchNorm1dNode(
                    new_layer_config(cfg.gnn.dim_inner, -1, -1, has_act=False,
                                     has_bias=False, cfg=cfg))
            # Update dim_in to reflect the new dimension fo the node features
            # 更新 dim_in 为经过编码后的维度 cfg.gnn.dim_inner
            self.dim_in = cfg.gnn.dim_inner
        # 是否使用边编码器
        if cfg.dataset.edge_encoder:
            # Hard-set edge dim for PNA.
            # 如果模型使用 PNA（Powerful Node Aggregation）层，将边特征的维度设为 16；否则使用 cfg.gnn.dim_inner。
            cfg.gnn.dim_edge = 16 if 'PNA' in cfg.gt.layer_type else cfg.gnn.dim_inner
            # 如果边编码器名称为 ER，则使用 EREdgeEncoder 进行编码。
            if cfg.dataset.edge_encoder_name == 'ER':
                self.edge_encoder = EREdgeEncoder(cfg.gnn.dim_edge)
            # 如果编码器名称以 +ER 结尾，使用额外的 ER 编码器处理部分边特征。
            elif cfg.dataset.edge_encoder_name.endswith('+ER'):
                # 动态加载配置中指定的边编码器。
                EdgeEncoder = register.edge_encoder_dict[
                    cfg.dataset.edge_encoder_name[:-3]]
                # 实例化边编码器。
                self.edge_encoder = EdgeEncoder(cfg.gnn.dim_edge - cfg.posenc_ERE.dim_pe)
                # 使用 ER 编码器处理边的附加属性。
                self.edge_encoder_er = EREdgeEncoder(cfg.posenc_ERE.dim_pe, use_edge_attr=True)
            else:
                EdgeEncoder = register.edge_encoder_dict[
                    cfg.dataset.edge_encoder_name]
                self.edge_encoder = EdgeEncoder(cfg.gnn.dim_edge)

            # 检查是否对边特征使用批归一化。
            if cfg.dataset.edge_encoder_bn:
                self.edge_encoder_bn = BatchNorm1dNode(
                    new_layer_config(cfg.gnn.dim_edge, -1, -1, has_act=False,
                                    has_bias=False, cfg=cfg))

    # 前向传播方法，用于将输入 batch 数据依次传入所有子模块。
    def forward(self, batch):
        # self.children() 方法仅返回模块中被注册为子模块的对象，即 torch.nn.Module 类的实例。只有当你在模型中定义了其他 nn.Module 对象时，它们才会被 children() 方法捕捉到并返回。
        # 遍历 FeatureEncoder 类中定义的所有子模块（如节点和边编码器）。
        for module in self.children():
            # 将 batch 数据传递给每个子模块进行处理。
            batch = module(batch)
        # 返回最终处理后的 batch。
        return batch


# 将 GPSModel 注册为 GraphGym 中的模型类型，方便后续通过名称引用。
@register_network('GPSModel')
class GPSModel(torch.nn.Module):
    """Multi-scale graph x-former.
    """

    def __init__(self, dim_in, dim_out):
        super().__init__()
        # 使用 FeatureEncoder 实例化特征编码器，输入特征维度为 dim_in。
        self.encoder = FeatureEncoder(dim_in)
        # 更新输入维度为编码后的维度。
        dim_in = self.encoder.dim_in

        # 如果配置中设置了 layers_pre_mp 的值大于 0，表示在主消息传递层之前有预处理层。
        if cfg.gnn.layers_pre_mp > 0:
            # 创建 GNNPreMP 实例，用于在图神经网络的主消息传递层之前执行预处理。
            self.pre_mp = GNNPreMP(
                dim_in, cfg.gnn.dim_inner, cfg.gnn.layers_pre_mp)
            # 更新输入维度为 cfg.gnn.dim_inner。
            dim_in = cfg.gnn.dim_inner

        # 断言语句，确保 cfg.gt.dim_hidden（隐藏层维度）、cfg.gnn.dim_inner（内部特征维度）和当前的 dim_in 一致。
        assert cfg.gt.dim_hidden == cfg.gnn.dim_inner == dim_in, \
            "The inner and hidden dims must match."

        # 尝试将 cfg.gt.layer_type 按 + 号分割为局部 GNN 类型和全局模型类型。如果格式不正确，抛出异常。
        try:
            local_gnn_type, global_model_type = cfg.gt.layer_type.split('+')
        except:
            raise ValueError(f"Unexpected layer type: {cfg.gt.layer_type}")
        layers = []
        # 根据配置中的 layers 参数，循环创建多个 GPSLayer 层。
        for _ in range(cfg.gt.layers):
            # 每一层都是 GPSLayer，并通过各种配置参数进行初始化，如隐藏层维度、GNN 类型、注意力机制的头数、PNA 度数、等价稳定 PE、dropout 等。
            layers.append(GPSLayer(
                dim_h=cfg.gt.dim_hidden,
                local_gnn_type=local_gnn_type,
                global_model_type=global_model_type,
                num_heads=cfg.gt.n_heads,
                pna_degrees=cfg.gt.pna_degrees,
                equivstable_pe=cfg.posenc_EquivStableLapPE.enable,
                dropout=cfg.gt.dropout,
                attn_dropout=cfg.gt.attn_dropout,
                layer_norm=cfg.gt.layer_norm,
                batch_norm=cfg.gt.batch_norm,
                bigbird_cfg=cfg.gt.bigbird,
            ))
        # 将所有层组成 Sequential 模块，按顺序执行。
        self.layers = torch.nn.Sequential(*layers)

        # 从 register.head_dict 中获取任务头（Head），用于处理最后的输出。
        GNNHead = register.head_dict[cfg.gnn.head]
        # 初始化任务头，输入特征维度为 cfg.gnn.dim_inner，输出维度为 dim_out。
        self.post_mp = GNNHead(dim_in=cfg.gnn.dim_inner, dim_out=dim_out)

    def forward(self, batch):
        for module in self.children():
            batch = module(batch)
        return batch
